package com.xp.ai.ragdemo;

import com.xp.ai.util.ApiKey;
import com.xp.ai.util.ModelUtils;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.DefaultContent;
import dev.langchain4j.rag.content.injector.ContentInjector;
import dev.langchain4j.rag.content.injector.DefaultContentInjector;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;

import java.util.ArrayList;
import java.util.List;

/**
 *
 * 测试一个完成的RAG匹配流程
 * @author xp
 */
public class RagMeituanDemo {

    public static void main(String[] args) {
        String question = "不想接单了，就拒绝接单了有处罚么？";
        EmbeddingStore<TextSegment> embeddingStore = PgVectorEmbeddingStore.builder()
                .host("localhost")
                .port(5433)
                .user("root")
                .password("123456")
                .database("postgres")
                .table("question_embedding")
                .dimension(1024)
                .useIndex(true)
                .indexListSize(100)
                .createTable(true)
                .dropTableFirst(false)
                .build();

        //创建向量化的模型
        OpenAiEmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .baseUrl(ApiKey.GJ_BASE_URL)
                .apiKey(ApiKey.GJ_API_KEY)
                .modelName(ApiKey.GJ_EMBEDDING_MODEL)
                .build();

        /*
        这一大段代码可以使用 ContentRetriever 来简化
        //1.首先将问题向量化
        Response<Embedding> embeddingResponse = embeddingModel.embed(question);
        //2.然后进行向量匹配
        EmbeddingSearchRequest request = EmbeddingSearchRequest.builder()
                .queryEmbedding(embeddingResponse.content())
                .minScore(0.0)
                .maxResults(5)
                .build();
        EmbeddingSearchResult<TextSegment> embeddingSearchResult = embeddingStore.search(request);
        List<EmbeddingMatch<TextSegment>> matches = embeddingSearchResult.matches();
        List<Content>  contents =  new ArrayList<>(100);
        for (EmbeddingMatch<TextSegment> match : matches) {
            DefaultContent content = new DefaultContent(match.embedded());
            contents.add(content);
        }*/

        //向量匹配搜索可以借助检索器
        ContentRetriever  contentRetriever =  new EmbeddingStoreContentRetriever(embeddingStore,embeddingModel,3);
        List<Content> contents = contentRetriever.retrieve(Query.from(question));


        //3.将匹配到的结果进行注入
        ContentInjector defaultContentInjector = new DefaultContentInjector();
        UserMessage injectUserMessage = defaultContentInjector.inject(contents, UserMessage.userMessage(question));
        System.out.println("inject = " + injectUserMessage);

        //4.将结果封装给大模型
        ChatLanguageModel chatLanguageModel = ModelUtils.getHuoshanR1Model();
        ChatResponse chatResponse = chatLanguageModel.chat(injectUserMessage);
        AiMessage aiMessage = chatResponse.aiMessage();
        String text = aiMessage.text();
        System.out.println(text);


    }
}
